#include "cybertron/network/SocketTcpServer.hpp"
#include "cybertron/network/SocketBase.hpp"
#include "cybertron/network/SocketEventHandler.hpp"
#include "cybertron/core/Log.hpp"
#include "cybertron/core/IdAllocator.hpp"
#include "cybertron/core/UtilTime.hpp"
#include "cybertron/core/UtilCRT.hpp"
#include "AsioSocketAsync.hpp"
#include <asio.hpp>
#include <thread>
#include <vector>
#include <set>

CYBERTRON_BEGIN

class SocketSet
{
public:
	int allocateId()
	{
		std::unique_lock<std::recursive_mutex> lock(mIdAllocatorMutex);
		return mIdAllocator.allocateId();
	}

	void freeId(int id)
	{
		if (id < 0)
		{
			return;
		}

		std::unique_lock<std::recursive_mutex> lock(mIdAllocatorMutex);
		mIdAllocator.releaseId(id);
	}

	void add(SocketBasePtr pSocket)
	{
		if (pSocket == nullptr)
		{
			return;
		}

		int id = allocateId();
		pSocket->setId(id);

		{
			std::unique_lock<std::recursive_mutex> lock(mSetMutex);
			mSet.insert(pSocket);
		}
	}

	void remove(SocketBasePtr pSocket)
	{
		if (pSocket == nullptr)
		{
			return;
		}

		int id = -1;

		{
			std::unique_lock<std::recursive_mutex> lock(mSetMutex);
			auto it = mSet.find(pSocket);
			if (it != mSet.end())
			{
				id = pSocket->getId();
				mSet.erase(it);
			}
		}

		if (id >= 0)
		{
			freeId(id);
		}
	}

	void clear()
	{
		std::unique_lock<std::recursive_mutex> lock(mSetMutex);
		mSet.clear();
	}

	size_t size() const
	{
		std::unique_lock<std::recursive_mutex> lock(mSetMutex);
		return mSet.size();
	}

	void getSockets(std::vector<std::shared_ptr<SocketBase>>& result)
	{
		std::vector<std::shared_ptr<SocketBase>> invalidSockets;

		{
			std::unique_lock<std::recursive_mutex> lock(mSetMutex);
			bool hasInvalidSocket = false;
			for (auto it = mSet.begin(); it != mSet.end(); ++it)
			{
				auto pSocket = (*it);
				if (pSocket->getStatus() != ESocketStatus_Connected)
				{
					hasInvalidSocket = true;
				}
				else
				{
					result.push_back(pSocket);
				}
			}

			if (hasInvalidSocket)
			{
				std::set<std::shared_ptr<SocketBase>> newSet;
				for (auto it = mSet.begin(); it != mSet.end(); ++it)
				{
					auto pSocket = (*it);
					if (pSocket->getStatus() == ESocketStatus_Connected)
					{
						newSet.insert(pSocket);
					}
					else
					{
						invalidSockets.push_back(pSocket);
					}
				}
				mSet.swap(newSet);
			}
		}

		if (invalidSockets.size() > 0)
		{
			for (size_t i = 0; i < invalidSockets.size(); ++i)
			{
				freeId(invalidSockets[i]->getId());
			}
		}
	}

private:
	mutable std::recursive_mutex mIdAllocatorMutex;
	IdAllocator mIdAllocator;

	mutable std::recursive_mutex mSetMutex;
	std::set<std::shared_ptr<SocketBase>> mSet;
};

class SocketEventHandlerDisconnect : public SocketEventHandler
{
public:
	SocketEventHandlerDisconnect(std::shared_ptr<SocketSet> pSocketSet)
	{
		mSocketSet = pSocketSet;
	}

	virtual void onDisconnected()
	{
		auto pSocket = getSocket();
		if (pSocket != nullptr)
		{
			logEvent("Disconnected from %s:%d", pSocket->getRemoteIp().c_str(), pSocket->getRemotePort());
		}
		if (mSocketSet != nullptr)
		{
			mSocketSet->remove(pSocket);
		}
	}

private:
	std::shared_ptr<SocketSet> mSocketSet;
};

class SocketTcpServerImpl
{
public:
	std::shared_ptr<asio::ip::tcp::acceptor> mpAcceptor;
	std::shared_ptr<asio::io_context> mpIoContext;
	std::vector<std::shared_ptr<std::thread>> mThreads;
	std::shared_ptr<SocketSet> mpSockets;
	std::atomic<unsigned int> mMaximalIdleMilliseconds;
	std::shared_ptr<std::thread> mClearIdleSocketThread;
	bool mRunning;
};

CybertronCoreAPI SocketTcpServer::SocketTcpServer()
{
	mpImpl = new SocketTcpServerImpl();

	mpImpl->mRunning = false;
	mpImpl->mMaximalIdleMilliseconds = 1000 * 30;
	mpImpl->mpSockets = std::make_shared<SocketSet>();
}

CybertronCoreAPI SocketTcpServer::~SocketTcpServer()
{
	closeThis();

	delete mpImpl;
	mpImpl = nullptr;
}

CybertronCoreAPI bool SocketTcpServer::listen(
	int port,
	int numThreads)
{
	close();

	mpImpl->mpIoContext = std::make_shared<asio::io_context>();

	asio::ip::tcp::resolver resolver(*mpImpl->mpIoContext);
	asio::ip::tcp::endpoint endpoints(asio::ip::tcp::v4(), port);
	mpImpl->mpAcceptor = std::make_shared<asio::ip::tcp::acceptor>(*(mpImpl->mpIoContext), endpoints);
	
	asio::ip::tcp::endpoint remote_ep = mpImpl->mpAcceptor->local_endpoint();
	unsigned short cur_listen_port = remote_ep.port();
	//asio::ip::address remote_ad = remote_ep.address();
	//	std::string s = remote_ad.to_string();
	logSilent("Listening to port %d...", cur_listen_port);
	doAccept();

	auto pContext = mpImpl->mpIoContext;
	unsigned int numCores = std::thread::hardware_concurrency();
	if (numThreads <= 0)
	{
		numThreads = (int)numCores;
	}
	for (int i = 0; i < numThreads; ++i)
	{
		auto pThread = UtilCRT::startThread(
			"Tcp Server IO", 
			[pContext] {
			if (pContext != nullptr)
			{
				pContext->run();
			}
		});

		mpImpl->mThreads.push_back(pThread);
	}

	mpImpl->mRunning = true;
	mpImpl->mClearIdleSocketThread = UtilCRT::startThread(
		"ClearIdle", 
		[this] {
		doClearIdle();
	});

	return true;
}

void SocketTcpServer::doAccept()
{
	auto self(shared_from_this());
	mpImpl->mpAcceptor->async_accept([self, this](asio::error_code ec, asio::ip::tcp::socket socket)
	{
		if (!ec)
		{
			socket.set_option(asio::ip::tcp::no_delay(true));
			asio::ip::address remoteAddr = socket.remote_endpoint().address();
			std::string remoteIp = remoteAddr.to_string();
			int remotePort = socket.remote_endpoint().port();

			logSilent("Accepted connection from %s:%d", remoteIp.c_str(), remotePort);

			std::shared_ptr<AsioSocketAsync> pSocketImpl = std::make_shared<AsioSocketAsync>(
				mpImpl->mpIoContext,
				std::move(socket),
				ESocketStatus_Connected,
				remoteIp,
				remotePort);

			std::shared_ptr<SocketBase> pSocket = std::make_shared<SocketBase>();
			pSocket->setImpl(pSocketImpl);

			pSocket->addHandler(
				std::make_shared<SocketEventHandlerDisconnect>(mpImpl->mpSockets));

			// add to list before onAccept to assign id to pSocket
			if (pSocket->getStatus() == ESocketStatus_Connected)
			{
				mpImpl->mpSockets->add(pSocket);
			}

			onAccept(pSocket);

			pSocketImpl->startReceiving();

		}
		else if(ec.value() == 24)
		{
			// Reach Linux open file number limitation, change it with ulimit -n 4096.
			logError("Socket Accept Error (%d): %s", ec.value(), ec.message().c_str());
			logError("Reach connection upper limitation, sleep 1 sec... If Linux, check: ulimit -n");
			std::this_thread::sleep_for(std::chrono::milliseconds(1000));
		}
		else
		{
			logError("Socket Accept Error (%d): %s", ec.value(), ec.message().c_str());
		}

		doAccept();
	});
}

void SocketTcpServer::doClearIdle()
{
	unsigned int tick = 0;
	unsigned int fps = 30;

	while (mpImpl->mRunning)
	{
		std::this_thread::sleep_for(std::chrono::milliseconds(1000 / fps));

		tick++;

		if (tick < fps)
		{
			continue;
		}

		tick = 0;

		unsigned int maximalIdleMilliseconds = mpImpl->mMaximalIdleMilliseconds;
		unsigned int currentTS = UtilTime::millisecondsCurrent();

		auto pSockets = mpImpl->mpSockets;
		if (pSockets != nullptr)
		{
			std::vector<std::shared_ptr<SocketBase>> sockets;
			sockets.reserve(pSockets->size());

			pSockets->getSockets(sockets);

			for (size_t i = 0; i < sockets.size(); ++i)
			{
				unsigned int lastMsgTS = sockets[i]->getLastMessageTime();
				if (lastMsgTS > currentTS)
				{
					// This could happend in multithreading
					continue;
				}
				unsigned int delta = currentTS - lastMsgTS;
				if (delta <= maximalIdleMilliseconds)
				{
					continue;
				}

				sockets[i]->close("Clear idle");
			}
		}
	}
}

CybertronCoreAPI void SocketTcpServer::close()
{
	closeThis();
}

CybertronCoreAPI size_t SocketTcpServer::getNumClients() const
{
	return mpImpl->mpSockets->size();
}

CybertronCoreAPI unsigned int SocketTcpServer::getMaximalIdleMilliseconds() const
{
	return mpImpl->mMaximalIdleMilliseconds;
}
CybertronCoreAPI unsigned short  SocketTcpServer::getListenPort() {
	asio::ip::tcp::endpoint remote_ep = mpImpl->mpAcceptor->local_endpoint();
	unsigned short cur_listen_port = remote_ep.port();
	return cur_listen_port;
}
CybertronCoreAPI void SocketTcpServer::setMaximalIdleMilliseconds(unsigned int maxIdleMilliseconds)
{
	mpImpl->mMaximalIdleMilliseconds = maxIdleMilliseconds;
}

CybertronCoreAPI void SocketTcpServer::broadcast(const Message& msg)
{
	if (mpImpl == nullptr)
	{
		return;
	}

	std::vector<std::shared_ptr<SocketBase>> sockets;
	sockets.reserve(mpImpl->mpSockets->size());

	mpImpl->mpSockets->getSockets(sockets);

	for (size_t i = 0; i < sockets.size(); ++i)
	{
		auto pSocket = sockets[i];
		Message msgClone = msg;
		pSocket->send(msgClone);
	}
}

CybertronCoreAPI void SocketTcpServer::broadcast(const std::string& msg)
{
    if (mpImpl == nullptr)
    {
        return;
    }

    std::vector<std::shared_ptr<SocketBase>> sockets;
    sockets.reserve(mpImpl->mpSockets->size());

    mpImpl->mpSockets->getSockets(sockets);

    for (size_t i = 0; i < sockets.size(); ++i)
    {
        auto pSocket = sockets[i];
        Message msgClone;
        msgClone.fromString(msg);
        pSocket->send(msgClone);
    }
}

void SocketTcpServer::closeThis()
{
	if (mpImpl != nullptr)
	{
		if (mpImpl->mpIoContext != nullptr)
		{
			mpImpl->mpIoContext->stop();
		}

		mpImpl->mpSockets->clear();

		for (size_t i = 0; i < mpImpl->mThreads.size(); ++i)
		{
			mpImpl->mThreads[i]->join();
		}
		mpImpl->mThreads.clear();

		mpImpl->mRunning = false;

		if (mpImpl->mClearIdleSocketThread != nullptr)
		{
			mpImpl->mClearIdleSocketThread->join();
			mpImpl->mClearIdleSocketThread = nullptr;
		}

		mpImpl->mpAcceptor = nullptr;
		mpImpl->mpSockets->clear();

		mpImpl->mpIoContext = nullptr;
	}
}

CybertronCoreAPI void SocketTcpServer::onAccept(SocketBasePtr pRemoteSocket)
{
}

CYBERTRON_END
